package jwtx

import (
	"context"
	"github.com/golang-jwt/jwt/v4"
	"github.com/zeromicro/x/errors"
	xhttp "github.com/zeromicro/x/http"
	"gorm.io/gorm"
	"net/http"
	"strings"
	"time"
)

type JwtAuthMiddleware struct {
	accessSecret string
	accessExpire int64
	db           *gorm.DB
}

func NewJwtAuthMiddleware(db *gorm.DB, accessSecret string, accessExpire int64) *JwtAuthMiddleware {
	return &JwtAuthMiddleware{
		db:           db,
		accessSecret: accessSecret,
		accessExpire: accessExpire,
	}
}

func (m *JwtAuthMiddleware) Handle(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		authHeader := r.Header.Get("Authorization")
		if authHeader == "" {
			xhttp.JsonBaseResponseCtx(r.Context(), w, errors.New(401, "Invalid Token"))
			return
		}

		// 兼容没有传Bearer
		var token string
		ss := strings.Split(authHeader, " ")
		if len(ss) == 2 {
			token = ss[1]
		} else if len(ss) == 1 {
			token = ss[0]
		}

		claims, err := m.ParseToken(token)
		if err != nil {
			xhttp.JsonBaseResponseCtx(r.Context(), w, errors.New(401, "Invalid Token"))
			return
		}

		ctx := r.Context()
		ctx = context.WithValue(ctx, "claims", claims)
		next.ServeHTTP(w, r.WithContext(ctx))

		//next(w, r.WithContext(ctx))
	}
}

func (m *JwtAuthMiddleware) Secret() jwt.Keyfunc {
	return func(token *jwt.Token) (interface{}, error) {
		return []byte(m.accessSecret), nil
	}
}

func (m *JwtAuthMiddleware) ParseToken(tokenStr string) (*CustomClaims, error) {
	token, err := jwt.ParseWithClaims(tokenStr, &CustomClaims{}, m.Secret())
	if err != nil {
		if ve, ok := err.(*jwt.ValidationError); ok {
			if ve.Errors&jwt.ValidationErrorMalformed != 0 {
				return nil, errors.New(403, "that's not even a token")
			} else if ve.Errors&jwt.ValidationErrorExpired != 0 {
				return nil, errors.New(403, "token is expired")
			} else if ve.Errors&jwt.ValidationErrorNotValidYet != 0 {
				return nil, errors.New(403, "token not active yet")
			} else {
				return nil, errors.New(403, "couldn't handle this token")
			}
		}
	}
	if claims, ok := token.Claims.(*CustomClaims); ok && token.Valid {
		return claims, nil
	}
	return nil, errors.New(403, "couldn't handle this token")
}

func (m *JwtAuthMiddleware) GenToken(uid int, roles []int, hasRootRole bool, actions map[int]int) (string, int64, error) {
	expireAt := jwt.NewNumericDate(time.Now().Add(time.Second * time.Duration(m.accessExpire)))
	claim := CustomClaims{
		Uid:         uid,
		Roles:       roles,
		HasRootRole: hasRootRole,
		Actions:     actions,
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: expireAt,                       //过期时间
			Issuer:    "JWT_ISSUER-EXAMPLE",           //签发人
			NotBefore: jwt.NewNumericDate(time.Now()), // 生效时间
		},
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claim)
	tokenS, err := token.SignedString([]byte(m.accessSecret))
	return tokenS, expireAt.Unix(), err
}
